#%%
import os
from pathlib import Path

print('WD:',os.getcwd())
# os.chdir()
PATH = Path().cwd()
print('PATH:',PATH)

#%%
'''
Computations (init).
Train loop (training_step)
Validation loop (validation_step)
Test loop (test_step)
Optimizers (configure_optimizers)
'''
# %%
# 最小化的代码模板
import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
import pytorch_lightning as pl


class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    # 上面的步骤和定义模型没什么区别

    # 定义训练步骤
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

    # 配置优化器
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

model = LitModel()
model
# %% 训练

train_loader = DataLoader(MNIST(os.getcwd(), download=True, transform=transforms.ToTensor()),batch_size=1000)
#%%
x,y = next(iter(train_loader))
print(x.shape)
print(y.shape)

#%%
trainer = pl.Trainer(max_epochs=3)
print(trainer.__dict__)